# -*- coding:utf-8 -*-
from __future__ import division
from __future__ import print_function
import sys
from pyspark import SparkContext
from pyspark import SQLContext, HiveContext

from pyspark.sql.functions import udf, broadcast
from pyspark.sql.types import *
from pyspark.ml import Pipeline
from pyspark.ml.regression import GBTRegressor, GBTRegressionModel
from pyspark.ml.classification import GBTClassifier, GBTClassificationModel
from pyspark.ml.feature import VectorIndexer
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.linalg import Vectors, VectorUDT

def get_prob(x):
    return float(x[0])


if __name__ == '__main__':

    model_file = 'hdfs:/user/wanghuaqiang/gender_model3'

    sc = SparkContext(appName="Train age predict model ")
    sqlContext = HiveContext(sc)

    feature_cols = [
        "avg_monthly_consum_amt",
        "avg_monthly_order_cnt",
        "last_3d_order_cnt",
        "last_7d_order_cnt",
        "last_15d_order_cnt",
        "last_30d_order_cnt",
        "last_90d_order_cnt",
        "last_180d_order_cnt",
        "last_3d_order_amt",
        "last_7d_order_amt",
        "last_15d_order_amt",
        "last_30d_order_amt",
        "last_90d_order_amt",
        "last_180d_order_amt",
        "last_3d_order_unit_amt",
        "last_7d_order_unit_amt",
        "last_15d_order_unit_amt",
        "last_30d_order_unit_amt",
        "last_90d_order_unit_amt",
        "last_180d_order_unit_amt",
        "is_last_3d_order",
        "is_last_7d_order",
        "is_last_15d_order",
        "is_last_30d_order",
        "is_last_90d_order",
        "is_last_180d_order",
        "last_3d_goods_cnt",
        "last_7d_goods_cnt",
        "last_15d_goods_cnt",
        "last_30d_goods_cnt",
        "last_90d_goods_cnt",
        "last_180d_goods_cnt",
        "last_3d_team_cnt",
        "last_7d_team_cnt",
        "last_15d_team_cnt",
        "last_30d_team_cnt",
        "last_90d_team_cnt",
        "last_180d_team_cnt",
        "last_3d_class1_cnt",
        "last_7d_class1_cnt",
        "last_15d_class1_cnt",
        "last_30d_class1_cnt",
        "last_90d_class1_cnt",
        "last_180d_class1_cnt",
        "last_3d_avg_daily_order_cnt",
        "last_7d_avg_daily_order_cnt",
        "last_15d_avg_daily_order_cnt",
        "last_30d_avg_daily_order_cnt",
        "last_90d_avg_daily_order_cnt",
        "last_180d_avg_daily_order_cnt",
        # "top3_ai_class1_list",
        "dndq_order_amt",
        "dndq_order_cnt",
        "dndq_goods_cnt",
        "dndq_team_cnt",
        "hfcz_order_amt",
        "hfcz_order_cnt",
        "hfcz_goods_cnt",
        "hfcz_team_cnt",
        "hnwy_order_amt",
        "hnwy_order_cnt",
        "hnwy_goods_cnt",
        "hnwy_team_cnt",
        "jjjz_order_amt",
        "jjjz_order_cnt",
        "jjjz_goods_cnt",
        "jjjz_team_cnt",
        "myyp_order_amt",
        "myyp_order_cnt",
        "myyp_goods_cnt",
        "myyp_team_cnt",
        "nnfs_order_amt",
        "nnfs_order_cnt",
        "nnfs_goods_cnt",
        "nnfs_team_cnt",
        "qt_order_amt",
        "qt_order_cnt",
        "qt_goods_cnt",
        "qt_team_cnt",
        "qcyp_order_amt",
        "qcyp_order_cnt",
        "qcyp_goods_cnt",
        "qcyp_team_cnt",
        "rybh_order_amt",
        "rybh_order_cnt",
        "rybh_goods_cnt",
        "rybh_team_cnt",
        "shfw_order_amt",
        "shfw_order_cnt",
        "shfw_goods_cnt",
        "shfw_team_cnt",
        "spcy_order_amt",
        "spcy_order_cnt",
        "spcy_goods_cnt",
        "spcy_team_cnt",
        "sjsm_order_amt",
        "sjsm_order_cnt",
        "sjsm_goods_cnt",
        "sjsm_team_cnt",
        "tsyx_order_amt",
        "tsyx_order_cnt",
        "tsyx_goods_cnt",
        "tsyx_team_cnt",
        "xlxb_order_amt",
        "xlxb_order_cnt",
        "xlxb_goods_cnt",
        "xlxb_team_cnt",
        "ydhw_order_amt",
        "ydhw_order_cnt",
        "ydhw_goods_cnt",
        "ydhw_team_cnt",
        "zbps_order_amt",
        "zbps_order_cnt",
        "zbps_goods_cnt",
        "zbps_team_cnt",
    ]

    udf_vectorize = udf(lambda *xs: Vectors.dense([float(x) for x in xs]), VectorUDT())
    udf_to_float = udf(lambda x: float(x), FloatType())
    udf_get_gender = udf(lambda x: int(x + 1), IntegerType())
    udf_get_prob = udf(lambda x: get_prob(x), FloatType())

    # 'select * from dw.dws_scrm_persona_customer_dim as a join dm_scrm.dm_persona_order_goods_index as b on a.yz_uid = b.buyer_id where (a.gender = 1 or a.gender=2) and last_180d_order_amt > 0 '
    # todo use only the goods table for eval
    sql = """
    select * from dm_scrm.dm_persona_customer_info
    """
    df = sqlContext.sql(sql)

    df = df.withColumn('features', udf_vectorize(*feature_cols)).withColumnRenamed('yz_uid', 'buyer_id')

    model = GBTClassificationModel.load(model_file)

    predictions = model.transform(df)
    cols = predictions.columns
    print(cols)
    # pred = predictions.select("buyer_id", "prediction", 'probability').withColumn()

    # print("预测\t标签")
    # for d in pred:
    #     p = d['prediction']
    #     t = d['gender']
    #     r = d['probability']
    #     print('\t'.join([str(k) for k in ([p, t, r])]))

    out_df = predictions.select('buyer_id', 'prediction', 'probability')\
        .withColumn('gender', udf_get_gender('prediction')).withColumn('prob', udf_get_prob('probability'))
    out_df.registerTempTable('prediction')
    sqlContext.sql("truncate table dm_ai.buyer_gender_predict_dev")  # for app user only
    sqlContext.sql("insert into table dm_ai.buyer_gender_predict_dev select buyer_id, gender, prob from prediction ")
    # out_df.write.mode('overwrite').saveAsTable('dm_ai.buyer_age_predict_dev')
    sc.stop()

